from openai import OpenAI


class KIMI:
    def __init__(self, api_key, base_url):
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url,
        )

    def __call__(self, *args, **kwargs):
        return self.generate(*args, **kwargs)

    def generate(self, model, prompt, system_prompt, temperature=0.3, stream=True):
        response = self.client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": f"{system_prompt}"},
                {"role": "user", "content": f"{prompt}"}
            ],
            temperature=temperature,
            stream=stream,
        )
        if stream:
            for chunk in response:
                delta = chunk.choices[0].delta
                if delta.content is not None:
                    yield delta.content
        else:
            for chunk in response.choices[0].message.content:
                yield chunk


if __name__ == "__main__":
    init_params = {
        "api_key": "sk-0cDK3fy39sQJ3FF8JCbgbDpepgpTSZrh6qyuhkqBbjZYCmHm",
        "base_url": "https://api.moonshot.cn/v1"
    }
    model = KIMI(**init_params)
    params = {
        "model": "moonshot-v1-8k",
        "prompt": "介绍一下自己",
        "system_prompt": "你是小艺，一个友善博学的AI助手",
        "temperature": 0.3,
        "stream": False,
    }

    res = model(**params)
    for i in res:
        print(i)
